# 绘制3D抛物线势阱中量子谐振子的本征态
# 例如，原子核的壳层模型认为核子处于原子核的势阱中
# 可能有bug
# 使用AI辅助
# Gitee Repo

import numpy as np
import scipy
import scipy.special
import matplotlib.pyplot as plt

L = 3;
dx = 0.1;

# 本征态的量子数，x,y,z方向
# nx,ny,nz=0,1,2,...
nx = 3;
ny = 3;
nz = 3;
n = nx+ny+nz #总的量子数，和能量有关 E= \hbar \omega(n+3/2)
coeff = 2; #势阱深度，理论上是 m\omega / (2\hbar^2)

x,y,z = np.meshgrid(np.arange(-L,L,dx),np.arange(-L,L,dx),np.arange(-L,L,dx),indexing='ij')
n = x.shape[0]

def compute_psi(x, n):
    hermiteH = scipy.special.hermite(n)
    psi = hermiteH(x)*np.exp(-coeff*x**2/2);
    return psi

psix = compute_psi(x,nx);
psiy = compute_psi(y,ny);
psiz = compute_psi(z,nz);
psi = psix*psiy*psiz;
u = psi**2 # 波函数平方是概率密度
u /= (np.sum(u)*dx**3) # 归一化

def draw_slice(x,y,z,u):
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
    from mpl_toolkits.mplot3d import Axes3D

    x_cpu = x
    y_cpu = y
    z_cpu = z
    u_cpu = u

    n = x_cpu.shape[0]

    plt.plot(x[:,int(n/2),int(n/2)], u_cpu[:,int(n/2),int(n/2)],
         marker='o',  # 数据点标记为圆圈
         linestyle='-',  # 实线连接
         color='blue',  # 线条颜色
         linewidth=2  # 线宽
         )
     
    # 添加标题和标签
    plt.title('x的psix切面')
    plt.xlabel('x', fontsize=12)
    plt.ylabel('u', fontsize=12)
    plt.axis([-L,L,-2,2])

    plt.show()
    return

def draw(x,y,z,u):
    import plotly.graph_objects as go
    from plotly.offline import plot

    downsample_factor = 2 #降低精度以避免过大的数据

    x_cpu = x[::downsample_factor, ::downsample_factor, ::downsample_factor]
    y_cpu = y[::downsample_factor, ::downsample_factor, ::downsample_factor]
    z_cpu = z[::downsample_factor, ::downsample_factor, ::downsample_factor]
    u_cpu = u[::downsample_factor, ::downsample_factor, ::downsample_factor]

    fig = go.Figure(data=go.Volume(
        x=x_cpu.flatten(),
        y=y_cpu.flatten(),
        z=z_cpu.flatten(),
        value=u_cpu.flatten(),
        cmin=0,
        cmax=2*np.median(u_cpu.flatten()),
        opacity=0.1, # needs to be small to see through all surfaces
        surface_count=20, # needs to be a large number for good volume rendering
        ))

    fig.update_layout(
        scene=dict(
            camera=dict(
                projection=dict(
                    type='orthographic'  # 设置为正交投影
                )
            )
        )
    )

    plot(fig)
    return

draw(x,y,z,u)
draw_slice(x,y,z,psix)
